// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <string>
#include <vector>

#include <kmock/kmock.h>
#include <ktest/ktest.h>

#include <nebula/testing/ktest_util.h>
#include <nebula/util/string.h>
#include <nebula/util/vector.h>

namespace nebula {
namespace internal {

TEST(StlUtilTest, VectorAddRemoveTest) {
  std::vector<int> values;
  std::vector<int> result = AddVectorElement(values, 0, 100);
  EXPECT_EQ(values.size(), 0);
  EXPECT_EQ(result.size(), 1);
  EXPECT_EQ(result[0], 100);

  // Add 200 at index 0 and 300 at the end.
  std::vector<int> result2 = AddVectorElement(result, 0, 200);
  result2 = AddVectorElement(result2, result2.size(), 300);
  EXPECT_EQ(result.size(), 1);
  EXPECT_EQ(result2.size(), 3);
  EXPECT_EQ(result2[0], 200);
  EXPECT_EQ(result2[1], 100);
  EXPECT_EQ(result2[2], 300);

  // Remove 100, 300, 200
  std::vector<int> result3 = DeleteVectorElement(result2, 1);
  EXPECT_EQ(result2.size(), 3);
  EXPECT_EQ(result3.size(), 2);
  EXPECT_EQ(result3[0], 200);
  EXPECT_EQ(result3[1], 300);

  result3 = DeleteVectorElement(result3, 1);
  EXPECT_EQ(result3.size(), 1);
  EXPECT_EQ(result3[0], 200);

  result3 = DeleteVectorElement(result3, 0);
  EXPECT_TRUE(result3.empty());
}

TEST(StlUtilTest, VectorFlatten) {
  std::vector<int> a{1, 2, 3};
  std::vector<int> b{4, 5, 6};
  std::vector<int> c{7, 8, 9};
  std::vector<std::vector<int>> vecs{a, b, c};
  auto actual = FlattenVectors(vecs);
  std::vector<int> expected{1, 2, 3, 4, 5, 6, 7, 8, 9};
  ASSERT_EQ(expected, actual);
}

TEST(StlUtilTest, VectorFilter) {
  std::vector<int> input{1, 2, 3, 4, 5, 6, 7, 8, 9};
  auto filtered = FilterVector(input, [](int i) { return i % 3 == 0; });

  EXPECT_THAT(filtered, ::testing::ElementsAre(3, 6, 9));
}

static std::string int_to_str(int val) { return std::to_string(val); }

TEST(StlUtilTest, VectorMap) {
  std::vector<int> input{1, 2, 3};
  std::vector<std::string> expected{"1", "2", "3"};

  auto actual = MapVector(int_to_str, input);
  ASSERT_EQ(expected, actual);

  auto bind_fn = std::bind(int_to_str, std::placeholders::_1);
  actual = MapVector(bind_fn, input);
  ASSERT_EQ(expected, actual);

  std::function<std::string(int)> std_fn = int_to_str;
  actual = MapVector(std_fn, input);
  ASSERT_EQ(expected, actual);

  actual = MapVector([](int val) { return std::to_string(val); }, input);
  ASSERT_EQ(expected, actual);
}

TEST(StlUtilTest, VectorMaybeMapFails) {
  std::vector<int> input{1, 2, 3};
  auto mapper = [](int item) -> turbo::Result<std::string> {
    if (item == 1) {
      return turbo::invalid_argument_error("XYZ");
    }
    return std::to_string(item);
  };
  ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, MaybeMapVector(mapper, input));
}

TEST(StlUtilTest, VectorMaybeMap) {
  std::vector<int> input{1, 2, 3};
  std::vector<std::string> expected{"1", "2", "3"};
  EXPECT_OK_AND_ASSIGN(
      auto actual,
      MaybeMapVector([](int item) -> turbo::Result<std::string> { return std::to_string(item); },
                     input));
  ASSERT_EQ(expected, actual);
}

TEST(StlUtilTest, VectorUnwrapOrRaise) {
  // TODO(ARROW-11998) There should be an easier way to construct these vectors
  std::vector<turbo::Result<MoveOnlyDataType>> all_good;
  all_good.push_back(turbo::Result<MoveOnlyDataType>(MoveOnlyDataType(1)));
  all_good.push_back(turbo::Result<MoveOnlyDataType>(MoveOnlyDataType(2)));
  all_good.push_back(turbo::Result<MoveOnlyDataType>(MoveOnlyDataType(3)));

  std::vector<turbo::Result<MoveOnlyDataType>> some_bad;
  some_bad.push_back(turbo::Result<MoveOnlyDataType>(MoveOnlyDataType(1)));
  some_bad.push_back(turbo::Result<MoveOnlyDataType>(turbo::invalid_argument_error("XYZ")));
  some_bad.push_back(turbo::Result<MoveOnlyDataType>(turbo::io_error("XYZ")));

  EXPECT_OK_AND_ASSIGN(auto unwrapped, UnwrapOrRaise(std::move(all_good)));
  std::vector<MoveOnlyDataType> expected;
  expected.emplace_back(1);
  expected.emplace_back(2);
  expected.emplace_back(3);

  ASSERT_EQ(expected, unwrapped);

  ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, UnwrapOrRaise(std::move(some_bad)));
}

}  // namespace internal
}  // namespace nebula
